import os.path

from lib.trainer_end_to_end import Trainer
from lib.config import *
from lib.dataset import SpeakerPairDataset
from lib.dataset_spliter import DatasetSpliter

fbank_feature_path = os.path.join(featureRoot, 'fbank')
dataset = SpeakerPairDataset(fbank_feature_path)
train_loader, test_loader = DatasetSpliter(dataset, trainPram['batchSize']).split()

# end_to_end_model_root = os.path.join(modelRoot, 'end_to_end_128')
# if not os.path.exists(end_to_end_model_root):
#     os.makedirs(end_to_end_model_root)
#
# trainer = Trainer(end_to_end_model_root, train_loader, test_loader, os.path.join(modelRoot, 'fbank', 'resnet34_custom', 'model'), preset='128')
# trainer.train()
#
# end_to_end_model_root = os.path.join(modelRoot, 'end_to_end_256')
# if not os.path.exists(end_to_end_model_root):
#     os.makedirs(end_to_end_model_root)
# trainer = Trainer(end_to_end_model_root, train_loader, test_loader, os.path.join(modelRoot, 'fbank', 'resnet34_custom', 'model'), preset='256')
# trainer.train()

end_to_end_model_root = os.path.join(modelRoot, 'end_to_end_512')
if not os.path.exists(end_to_end_model_root):
    os.makedirs(end_to_end_model_root)
trainer = Trainer(end_to_end_model_root, train_loader, test_loader, os.path.join(modelRoot, 'fbank', 'resnet34_custom', 'model'), preset='512')
trainer.train()

# end_to_end_model_root = os.path.join(modelRoot, 'end_to_end_1024')
# if not os.path.exists(end_to_end_model_root):
#     os.makedirs(end_to_end_model_root)
# trainer = Trainer(end_to_end_model_root, train_loader, test_loader, os.path.join(modelRoot, 'fbank', 'resnet34_custom', 'model'), preset='1024')
# trainer.train()
